#include "point_light_system.hpp"

#include <iostream>
#include <stdexcept>
#include <array>
#include <map>

#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#include <glm/glm.hpp>
#include <glm/gtc/constants.hpp>


namespace njm {

    struct PointLightPushConstants
    {
        glm::vec4 position{};
        glm::vec4 color{};
        float radius;
    };
    
    PointLightSystem::PointLightSystem(NjmDevice &device, VkRenderPass renderPass, VkDescriptorSetLayout globalSetLayout) : njmDevice{device}
    {
        createPipelineLayout(globalSetLayout);
        createPipeline(renderPass);
    }
    PointLightSystem::~PointLightSystem()
    {
        vkDestroyPipelineLayout(njmDevice.device(), pipelineLayout, nullptr);
    }

    void PointLightSystem::createPipelineLayout(VkDescriptorSetLayout globalSetLayout)
    {
        VkPushConstantRange pcRange{};
        pcRange.stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT;
        pcRange.offset = 0;
        pcRange.size = sizeof(PointLightPushConstants);

        std::vector<VkDescriptorSetLayout> setLayouts{globalSetLayout};

        VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
        pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
        pipelineLayoutInfo.setLayoutCount = static_cast<uint32_t>(setLayouts.size());
        pipelineLayoutInfo.pSetLayouts = setLayouts.data();
        pipelineLayoutInfo.pushConstantRangeCount = 1;
        pipelineLayoutInfo.pPushConstantRanges = &pcRange;

        if (vkCreatePipelineLayout(njmDevice.device(), &pipelineLayoutInfo, nullptr, &pipelineLayout) != VK_SUCCESS)
        {throw std::runtime_error("PointLightSystem::createPipelineLayout : failed to create pipelineLayout");}      
    }

    void PointLightSystem::createPipeline(VkRenderPass renderPass)
    {
        assert(pipelineLayout != nullptr && "PointLightSystem::createPipeline : pipelineLayout is nullptr");
        PipelineConfigInfo pipelineConfig{};
        NjmPipeline::defaultPipelineConfigInfo(pipelineConfig);
        NjmPipeline::enableAlphaBlending(pipelineConfig);
        pipelineConfig.renderPass = renderPass;
        pipelineConfig.pipelineLayout = pipelineLayout;
        pipelineConfig.bindingDescriptions.clear();
        pipelineConfig.attributeDescriptions.clear();

        njmPipeline = std::make_unique<NjmPipeline>(
            njmDevice, 
            "src/shaders/point_light.vert.spv", 
            "src/shaders/point_light.frag.spv",
            pipelineConfig);
    }

    void PointLightSystem::render(FrameInfo &frameInfo)
    {
        // sort light by dis^2
        glm::vec3 cameraPos = frameInfo.camera.getPosition();
        std::map<float, NjmGameObj::id_t> sortedLights;
        for (auto &kv : frameInfo.gameObjects)
        {
            auto& obj = kv.second;
            if(obj.pointLight == nullptr) continue;
            glm::vec3 offset = obj.transform.translation - cameraPos;
            float disSquared = glm::dot(offset, offset);
            sortedLights[disSquared] = obj.getId();
        }

        njmPipeline->bind(frameInfo.commandBuffer);
        vkCmdBindDescriptorSets(frameInfo.commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
            pipelineLayout, 0, 1, &frameInfo.globalDescriptorSet, 0, nullptr);
        // vkCmdDraw(frameInfo.commandBuffer, 6, 1, 0, 0);

        uint32_t idx = 0;
        for (auto it = sortedLights.rbegin(); it != sortedLights.rend(); it++)
        {
            if(idx >= MAX_LIGHT_COUNT) break;
            auto& obj = frameInfo.gameObjects.at(it->second);
            // if(obj.pointLight == nullptr) continue;

            PointLightPushConstants pc{};
            pc.position = glm::vec4(obj.transform.translation, 1.0f);
            pc.color = glm::vec4(obj.color, obj.pointLight->lightIntensity);
            pc.radius = obj.transform.scale.x;
            vkCmdPushConstants(frameInfo.commandBuffer, pipelineLayout, 
                VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 
                0, sizeof(PointLightPushConstants), &pc);

            vkCmdDraw(frameInfo.commandBuffer, 6, 1, 0, 0);    
            idx += 1;
        }
    }

    void PointLightSystem::update(FrameInfo &frameInfo, GlobalUbo &ubo)
    {
        auto rotation = glm::rotate(glm::mat4(1.0f), frameInfo.frameTime, {0.0f, -1.0f, 0.0f});
        auto scale = 0.1f * glm::abs(glm::sin(0.1f * frameInfo.frameTime)) + 0.1f;
        auto selfRotation = glm::vec3(glm::sin(frameInfo.frameTime));

        uint32_t idx = 0;
        for (auto &kv : frameInfo.gameObjects)
        {
            if(idx >= MAX_LIGHT_COUNT) break;
            auto& obj = kv.second;
            if(obj.pointLight == nullptr) continue;

            obj.transform.translation = glm::vec3(rotation * glm::vec4(obj.transform.translation, 1.0f));
            obj.transform.scale = glm::vec3(scale);
            obj.transform.rotation = selfRotation;

            ubo.pointLights[idx].position = glm::vec4(obj.transform.translation, 1.0f);
            ubo.pointLights[idx].color = glm::vec4(obj.color, obj.pointLight->lightIntensity);
            idx += 1;
        }
        ubo.numLights = idx;
    }
}